﻿#pragma once
#include "Common.h"
#include "libmesh/enum_order.h"
#include <vtkm/CellShape.h>
#include <vtkm/Matrix.h>

using Point = Vec3f;
using Gradient = Vec3f;
using Matrix22 = vtkm::Matrix<FloatDefault, 2, 2>;
using Matrix33 = vtkm::Matrix<FloatDefault, 3, 3>;

template<typename Shape, libMesh::Order order>
struct Lagrange
{
};

template<>
struct Lagrange<vtkm::CellShapeTagTriangle, libMesh::Order::FIRST>
{
  static constexpr int NUM_SHAPES = 3;

  VTKM_EXEC static inline Vec<FloatDefault, NUM_SHAPES> shape(const Point& local)
  {
    auto xi = local[0];
    auto eta = local[1];
    return { xi, eta, 1 - xi - eta };
  }
  VTKM_EXEC static inline Vec<Gradient, NUM_SHAPES> dShapeAsXi(const Point& local)
  {
    return { { 1, 0, 0 }, { 0, 1, 0 }, { -1, -1, 0 } };
  }

  template<typename PointVector>
  VTKM_EXEC static inline Gradient derivativeShape(const PointVector& pts,
                                                   const Point& local,
                                                   int i);
  template<typename PointVector>
  VTKM_EXEC static inline Matrix22 jacobian(const PointVector& pts, const Point& local)
  {
    auto dphi_dxi = dShapeAsXi(local);

    Matrix22 m(0);
    for (size_t i = 0; i < NUM_SHAPES; i++)
    {
      for (size_t p = 0; p < 2; p++)
      {
        for (size_t q = 0; q < 2; q++)
        {
          m(p, q) += pts[i][p] * dphi_dxi[i][q];
        }
      }
    }

    return m;
  }
};

template<>
struct Lagrange<vtkm::CellShapeTagLine, libMesh::Order::FIRST>
{
  static constexpr int NUM_SHAPES = 2;

  VTKM_EXEC static inline Vec<FloatDefault, NUM_SHAPES> shape(FloatDefault& xi)
  {
    return { FloatDefault(0.5) * (1 - xi), FloatDefault(0.5) * (1 + xi) };
  }

  VTKM_EXEC static inline Vec<FloatDefault, NUM_SHAPES> shape(const Point& local)
  {
    FloatDefault xi = local[0];
    return shape(xi);
  }

  VTKM_EXEC static inline Vec<Gradient, NUM_SHAPES> dShapeAsXi(const Point& local)
  {
    return { { -0.5, 0, 0 }, { 0.5, 0, 0 } };
  }
};

template<>
struct Lagrange<vtkm::CellShapeTagQuad, libMesh::Order::FIRST>
{
  static constexpr int NUM_SHAPES = 4;
  using FELine = Lagrange<vtkm::CellShapeTagLine, libMesh::Order::FIRST>;

  VTKM_EXEC static inline Vec<FloatDefault, NUM_SHAPES> shape(const Point& local)
  {
    static const UInt8 i0[] = { 0, 1, 1, 0 };
    static const UInt8 i1[] = { 0, 0, 1, 1 };

    auto xi = local[0];
    auto eta = local[1];
    auto shape_line_xi = FELine::shape(xi);
    auto shape_line_eta = FELine::shape(eta);
    return { shape_line_xi[i0[0]] * shape_line_eta[i1[0]],
             shape_line_xi[i0[1]] * shape_line_eta[i1[1]],
             shape_line_xi[i0[2]] * shape_line_eta[i1[2]],
             shape_line_xi[i0[3]] * shape_line_eta[i1[3]] };
  }

  VTKM_EXEC static inline Vec<Gradient, NUM_SHAPES> dShapeAsXi(const Point& local)
  {
    auto dshape = FELine::dShapeAsXi(local);
    auto xi = local[0];
    auto eta = local[1];
    auto shape_line_xi = FELine::shape(xi);
    auto shape_line_eta = FELine::shape(eta);
    static const UInt8 i0[] = { 0, 1, 1, 0 };
    static const UInt8 i1[] = { 0, 0, 1, 1 };

    return {
      { dshape[i0[0]][0] * shape_line_eta[i1[0]], shape_line_eta[i0[0]] * dshape[i1[0]][0], 0 },
      { dshape[i0[1]][0] * shape_line_eta[i1[1]], shape_line_eta[i0[1]] * dshape[i1[1]][0], 0 },
      { dshape[i0[2]][0] * shape_line_eta[i1[2]], shape_line_eta[i0[2]] * dshape[i1[2]][0], 0 },
      { dshape[i0[3]][0] * shape_line_eta[i1[3]], shape_line_eta[i0[3]] * dshape[i1[3]][0], 0 }
    };

  }

  template<typename PointVector>
  VTKM_EXEC static inline Matrix22 jacobian(const PointVector& pts, const Point& local)
  {
    Vec<Vec3f, NUM_SHAPES> dphi = dShapeAsXi(local);
    Matrix22 m(0);
    for (size_t i = 0; i < NUM_SHAPES; i++)
    {
      for (size_t p = 0; p < 2; p++)
      {
        for (size_t q = 0; q < 2; q++)
        {
          m(p, q) += pts[i][p] * dphi[i][q];
        }
      }
    }

    //std::cout << m << std::endl;
    return m;
  }
};